import cv2
import json

import tqdm

from config import raw_dataset_path, generate_dataset_images_path, generate_dataset_label_path
import os
import numpy as np


def _exit():
    exit()


class RawDataDir:
    Labels = {
        "Legend": 0,
        "ValueTitle": 1,
        "Title": 2,
        "CategoryTitle": 3,
        "BarCanvas": 4,
        "LineCanvas": 5,
        "PieCanvas": 6,
        "Bar": 7,
        "Line": 8,
        "Pie": 9,
        "Pictogram": 10
    }

    def __init__(self, raw_dataset_dir, _sub_dir):
        self.raw_dataset = os.path.join(raw_dataset_dir, _sub_dir)
        self.raw_labels_dir = os.path.join(self.raw_dataset, "pictograms_labels")
        self.sub_dir = _sub_dir
        self.raw_labels = dict()
        self.middle_labels = dict()
        self.generative_labels = dict()
        self.load_labels()
        self.check_labels()

    def load_labels(self):
        """
        加载原始labels为字典，字典key为文件名，value为shape
        """
        # 判断是否存在
        if not os.path.exists(self.raw_labels_dir):
            return

        for label_filename in os.listdir(self.raw_labels_dir):
            label_filepath = os.path.join(self.raw_labels_dir, label_filename)
            with open(label_filepath, "r") as _fp:
                self.raw_labels[label_filename] = json.load(_fp)["shapes"]

    @staticmethod
    def check_include(keypoint_rectangle_list):
        def check_2_keypoint_rectangle_include(rec1, rec2):
            rec1_c = [sum(rec1[::2]) / 2, sum(rec1[1::2]) / 2]
            rec2_c = [sum(rec2[::2]) / 2, sum(rec2[1::2]) / 2]

            if rec2[0] < rec1_c[0] < rec2[2] and rec2[1] < rec1_c[1] < rec2[3]:
                return True

            if rec1[0] < rec2_c[0] < rec1[2] and rec1[1] < rec2_c[1] < rec1[3]:
                return True

            return False

        length = len(keypoint_rectangle_list)
        for i in range(length):
            for j in range(i + 1, length):
                if check_2_keypoint_rectangle_include(keypoint_rectangle_list[i]["points"],
                                                      keypoint_rectangle_list[j]["points"]):
                    return True
        return False

    def check_labels(self):
        """
        1、将labels标注转换格式
            文件名：[
                {
                    pictogram: [],
                    以及所有pictogram标签内的所有标签
                }
            ]
        2、检查labels标注是否存在问题，基于以下规则：
        - 每一个pictogram内不存在pictogram标签
        - 每一个pictogram标签内只存在一个canvas标签
        - 每一个pictogram内最多只存在一种数据标签
        - 每一个pictogram内各类title标签最多只存在一个，并且最多只有一个图例


        下面是shapes的类型定义：
        shapes : [{
        "label": "Pictogram",
          "points": [
            [
              254.11111111111114,
              690.1234567901234
            ],
            [
              1530.6543209876543,
              1522.2222222222222
            ]
            ]}]
        """
        for key in self.raw_labels.keys():
            shapes = self.raw_labels[key]
            pictogram_labels = []
            other_labels = []
            for shape in shapes:
                # 首先检查是否存在未定义标签
                if shape["label"] not in self.Labels.keys():
                    label = shape["label"]
                    print(f"In {self.raw_labels_dir}/{key}, there is a undefined label named {label}!")
                    _exit()

                # 展平
                new_shape = {
                    "label": shape["label"],
                    "points": sum(shape["points"], [])
                }

                if shape["label"] == "Pictogram":
                    pictogram_labels.append(new_shape)
                else:
                    other_labels.append(new_shape)
            # 检查pictogram_labels是否互相包含
            if self.check_include(pictogram_labels):
                print(f"In {self.raw_labels_dir}/{key}, there are some pictogram labels include each other!")
                _exit()

            # 将同一幅统计图的标注添加到一起
            pictograms = []
            for pictogram_label in pictogram_labels:
                tmp = {"Pictogram": pictogram_label["points"]}
                rec = pictogram_label["points"]
                li = []
                pictogram_type = 0
                for other_label in other_labels:
                    other_points = other_label["points"]
                    if rec[0] < other_points[0] < rec[2] and rec[1] < other_points[1] < rec[3]:
                        label_id = self.Labels[other_label["label"]]

                        # title 标签和 legend标签是否重复
                        if label_id in [0, 1, 2, 3] and label_id in li:
                            print(label_id)
                            print(f"In {self.raw_labels_dir}/{key}, "
                                  f"there is a single label (legend title) but appear twice in a pictogram!")
                            _exit()

                        if label_id in [4, 5, 6]:
                            if label_id in li:
                                print(f"In {self.raw_labels_dir}/{key}, "
                                      f"there is a single label (canvas) but appear twice in a pictogram!")
                            else:
                                pictogram_type = label_id

                        li.append(label_id)
                        if label_id in [7, 8, 9]:
                            if other_label["label"] not in tmp.keys():
                                tmp[other_label["label"]] = [other_label["points"]]
                            else:
                                tmp[other_label["label"]].append(other_label["points"])
                            continue
                        else:
                            tmp[other_label["label"]] = other_label["points"]

                # 没有检测到绘图区域
                if pictogram_type == 0:
                    print(f"In {self.raw_labels_dir}/{key}, there is not a canvas!")
                    _exit()

                data_element_labels = set([i for i in li if i > 6])
                if len(data_element_labels) > 1:  # 判断是否存在多个数据元素
                    print(f"In {self.raw_labels_dir}/{key}, there is a single label but appear twice in a pictogram!")
                    return
                else:  # 判断数据元素和canvas是否匹配
                    if len(data_element_labels) == 1:
                        data_element_label = list(data_element_labels)[0]
                        if data_element_label - 3 != pictogram_type:
                            print(f"In {self.raw_labels_dir}/{key}, canvas type is not equal to data element type!")
                            _exit()

                pictograms.append(tmp)

            self.middle_labels[key] = pictograms

    def generate_final_data(self):
        """
        生成最后的结果
        图片输出到
        名字为 {sub_dir}_{filename_prefix}_{index}.png
        label为[
                {sub_dir}_{filename_prefix}_{index}.png: {
                    has_data_element: True or False,
                    data_type: 0 or 1 or 2 ,
                    9个类别标注
                }
            ]
        self.middle_labels = {
            "{filename}": [
                {
                    "Pictogram": [x1, y1, x2, y2],
                    "Title": [x1, y1, x2, y2],
                    ...
                }
            ]
        }
        """
        for filename in self.middle_labels:
            tmp_index = 0
            filename_prefix = filename.split(".")[0]
            filepath = os.path.join(self.raw_dataset, filename_prefix + '.png')
            img = cv2.imdecode(np.fromfile(filepath, dtype=np.uint8), -1)
            for pictogram in self.middle_labels[filename]:
                cut_image_name = f"{self.sub_dir}_{filename_prefix}_{tmp_index}.png"
                tmp_index += 1
                cut_image_save_path = os.path.join(generate_dataset_images_path, cut_image_name)
                cip = pictogram["Pictogram"]
                pictogram_img = img[int(cip[1]):int(cip[3]), int(cip[0]): int(cip[2])]
                # 保存图片
                cv2.imwrite(cut_image_save_path, pictogram_img, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])
                # 下面保存标注
                base_x = cip[0]
                base_y = cip[1]
                has_data_element = False
                data_type = -1

                labels = dict()
                for label_type in pictogram.keys():
                    if label_type == "Pictogram":
                        continue

                    if self.Labels[label_type] > 6:
                        has_data_element = True

                    if 3 < self.Labels[label_type] < 7:
                        data_type = self.Labels[label_type] - 4

                    flatten_points = pictogram[label_type]
                    transform_points = []
                    if isinstance(flatten_points[0], list):  # 数据元素
                        for flatten_points_item in flatten_points:
                            transform_points_item = []
                            for i in range(len(flatten_points_item)):
                                if i % 2 == 0:  # X
                                    transform_points_item.append(flatten_points_item[i] - base_x)
                                else:  # Y
                                    transform_points_item.append(flatten_points_item[i] - base_y)
                            transform_points.append(transform_points_item)
                    else:
                        for i in range(len(flatten_points)):
                            if i % 2 == 0:  # X
                                transform_points.append(flatten_points[i] - base_x)
                            else:  # Y
                                transform_points.append(flatten_points[i] - base_y)

                    labels[label_type] = transform_points
                    labels["data_type"] = data_type
                    labels["has_data_element"] = has_data_element
                self.generative_labels[cut_image_name] = labels


if __name__ == '__main__':
    sub_dirs = os.listdir(raw_dataset_path)
    labels = []
    for sub_dir in tqdm.tqdm(sub_dirs):
        data = RawDataDir(raw_dataset_path, sub_dir)
        data.generate_final_data()
        labels.append(data.generative_labels)

    labels_map = dict()
    for label in labels:
        for key in label.keys():
            labels_map[key] = label[key]

    with open(generate_dataset_label_path, "w") as fp:
        json.dump(labels_map, fp)

    print("Finish!")
